# -*- coding: utf-8 -*- 
# Author: Jacky
# Creation Date: 2021/3/16
import numpy as np
import pandas as pd
import copy
import torch
from data_set import DataSet
from user_info import UserInfo
from options import args_parser
from data_set import convert
import utils
from options import args_parser
# 用来判断使用哪一种参数的
args = args_parser()


def get_dataset(args):
    """
    :param: args:
    :return: ratings: dataFrame ['user_id' 'movie_id' 'rating']
    :return: user_info:  dataFrame ['user_id' 'gender' 'age' 'occupation']
    """
    model_manager = utils.ModelManager('data_set') # 此方法中path_name和save_name的区别，'dataset'决定了路径model/data_set，后续模型的save和Load都是从这个路径
    user_manager = utils.UserInfoManager(args.dataset) # args.dataset默认为ml-100k，args.dataset决定了路径user/，后续

    '''Do you want to clean workspace and retrain model/data_set user again?'''
    '''if you want to retrain model/data_set user, please set clean_workspace True'''
    model_manager.clean_workspace(args.clean_dataset) # args.clean_dataset默认为false,删除model/dataset-100k下所有文件
    user_manager.clean_workspace(args.clean_user) # args.clean_user默认为false，删除user/下所有文件

    # 导入模型信息，模型导入文件
    try:
        ratings = model_manager.load_model(args.dataset + '-ratings') # 从model/data_set/ml-100k-ratings.pkl加载模型，args.dataset对应ml-100k
        print("Load " + args.dataset + " data_set success.\n")
    except OSError:
        ratings = DataSet.LoadDataSet(name=args.dataset) # 返回数据集ml-100k处理后的dataframe，根据user_id排序，去掉time_stamp。ratings就是model，dataframe格式
        model_manager.save_model(ratings, args.dataset + '-ratings') # 再把ratings（dataframe）转换为二进制文件，即ml-100k-ratings.pkl,模型的存储

    # 导入用户信息
    try:
        user_info = user_manager.load_user_info('user_info')
        print("Load " + args.dataset + " user_info success.\n")
    except OSError:
        user_info = UserInfo.load_user_info(name=args.dataset)
        user_manager.save_user_info(user_info, 'user_info') # user/ml-100k-user_info.csv
    # 返回两个dataframe对象，也就是model，分别对应评分集和用户集（不是convert之后的）
    return ratings, user_info

# sampling函数：1.加载或生成客户端数据集信息：首先，函数尝试从磁盘中加载已有的客户端数据集信息，包括训练集和测试集的索引以及样本数据。如果不存在这些信息，则会根据输入的参数重新生成。2. 划分数据集：将总体数据集分成训练集和测试集，并为每个客户端生成对应的训练集和测试集。这里采用了随机划分的方式，将每个客户端80%的数据作为训练集，剩余20%作为测试集。3.存储客户端数据集信息：生成或加载完成客户端数据集信息后，将其存储到磁盘上，以便后续实验使用。4.返回结果：返回样本数据集以及每个客户端的训练集和测试集索引。
# sample会对ratings和user_info进行内连接，生成新的dataframe对象
def sampling(args):
    """
    :param args
    :return: sample: matrix user_id|movie_id|rating|gender|age|occupation|label
    :return: user_group_train, the idx of sample for each client for training
    :return: user_group_test, the idx of sample for each client for testing
    """
    # 存储每个client信息
    model_manager = utils.ModelManager('clients') # model/clients
    '''Do you want to clean workspace and retrain model/clients again?'''
    '''if you want to change test_size or retrain model/clients, please set clean_workspace True'''
    model_manager.clean_workspace(args.clean_clients)
    # 导入模型信息
    # 如果存在已有的客户端信息，则直接加载；否则调用get_dataset(args)函数获取整体数据集的评分信息和用户信息，并根据这些信息生成每个客户端的训练集和测试集。
    try:
        if args.clients_num == 10:
            users_group_train = model_manager.load_model(args.dataset + '-user_group_train10') # model/clients/ml-100k-user_group_train10.pkl
            users_group_test = model_manager.load_model(args.dataset + '-user_group_test10') # model/clients/ml-100k-user_group_test10.pkl
        if args.clients_num == 100:
            users_group_train = model_manager.load_model(args.dataset + '-user_group_train100')
            users_group_test = model_manager.load_model(args.dataset + '-user_group_test100')
        sample = model_manager.load_model(args.dataset + '-sample') # 如ml-100k-sample.pkl,反序列化
        print("Load " + args.dataset + " clients info success.\n")
    except OSError:
        # 调用get_dataset函数，得到ratings,user_info
        ratings, user_info = get_dataset(args)
        # 每个client包含的用户数
        users_num_client = int((user_info.index[-1] + 1) / args.clients_num) # args.clients_num默认为10
        # sample user_id|movie_id|rating|gender|age|occupation
        # 将两个数据框（DataFrame）ratings 和 user_info 根据共同的列 user_id 进行内连接（inner join）
        sample = pd.merge(ratings, user_info, on=['user_id'], how='inner')
        # 使用 astype() 方法将 sample 中的列转换为指定的数据类型。其中，字典的键是列名，值是要转换的目标数据类型
        # 之所以能可以用user_id等索引是因为dataset.py中指定了pd.read_csv中names=ratings_header
        sample = sample.astype({'user_id': 'int64', 'movie_id': 'int64', 'rating': 'float64',
                                'gender': 'float64', 'age': 'float64', 'occupation': 'float64'})
        # 生成每个客户用来train和test的idx
        users_group_all, users_group_train, users_group_test = {}, {}, {}
        # 生成用户数据集
        # 对用户数据集进行划分train/test
        for i in range(args.clients_num):
            print('loading client ' + str(i))
            index_begin = ratings[ratings['user_id'] == int(users_num_client) * i + 1].index[0] # 选择了符合条件的第一行的索引，即客户端在样本数据集中的起始索引位置，返回满足条件的第一个元素的索引值（因为索引用户有多个）
            index_end = ratings[ratings['user_id'] == users_num_client * (i + 1)].index[-1] \
                if i != args.clients_num-1 else ratings.index[-1]
            # 创建了一个从 index_begin 到 index_end （包括 index_begin，但不包括 index_end + 1）的整数序列
            users_group_all[i] = set(np.arange(index_begin, index_end + 1))
            NUM_train = int(0.8 * len(users_group_all[i]))
            # 从 users_group_all[i] 中随机选择了 NUM_train 个索引作为训练集的索引
            users_group_train[i] = set(np.random.choice(list(users_group_all[i]), NUM_train, replace=False))
            # 剩下的为测试集的索引
            users_group_test[i] = users_group_all[i] - users_group_train[i]
            # 将set转换回list，并排序
            users_group_train[i] = list(users_group_train[i])
            users_group_test[i] = list(users_group_test[i])
            users_group_train[i].sort()
            users_group_test[i].sort()
            print('generate client ' + str(i) + ' info success\n')
        # 存储user_group_train user_group_test sample
        model_manager.save_model(sample, args.dataset + '-sample')
        if args.clients_num == 10:
            model_manager.save_model(users_group_train, args.dataset + '-user_group_train10')
            model_manager.save_model(users_group_test, args.dataset + '-user_group_test10')
        if args.clients_num == 100:
            model_manager.save_model(users_group_train, args.dataset + '-user_group_train100')
            model_manager.save_model(users_group_test, args.dataset + '-user_group_test100')
    # 返回样本数据集以及每个客户端的训练集和测试集索引
    # sample是从.pkl文件，如ml-100k-sample.pkl中加载的python对象，users_group_train, users_group_test是一维训练集和测试集用户索引，均为列表
    return sample, users_group_train, users_group_test


# w 是一个列表，其中每个元素都是一个字典，表示不同模型的权重，返回权重平均值
def average_weights(w):
    """
    Returns the average of the weights.
    """
    w_avg = copy.deepcopy(w[0])
    for key in w_avg.keys():
        for i in range(1, len(w)):
            w_avg[key] += w[i][key]
        w_avg[key] = torch.div(w_avg[key], len(w))
    return w_avg

# 随机找一个
if __name__ == '__main__':
    args = args_parser()
    ratings, user_info = get_dataset(args)
    sample, users_group_train, users_group_test = sampling(args)
    # 验证convert
    # 使用索引 6 从测试集中获取第6个客户端的样本数据，并将其转换为NumPy数组
    client_6 = np.array(sample.iloc[users_group_test[6], :])
    # 调用了一个名为 convert 的函数，将第六个客户端的样本数据转换为一定格式的用户-电影矩阵，并将结果赋值给 user_movie_6 变量
    user_movie_6 = convert(client_6, max(sample['movie_id'])) # 对应convert函数中client_6必须为np数组
    print(sample.head(50))# 经过实测sample中的数据是根据user_id有序的
    # print(users_group_train)
    # print(users_group_test)
    # print(user_movie_6)
